'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 3_study_cache_memory_with_buffer.py
* @Time: 2025/7/20
* @All Rights Reserve By Brtc
'''
from operator import itemgetter

import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate ,MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferWindowMemory

dotenv.load_dotenv()
#1、创建提示词模板& 记忆
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是OpenAI 开发的机器人，请根据对应的上下文来回答用户的问题。"),
    MessagesPlaceholder("history"),
    ("human", "{query}"),
])
memory = ConversationBufferWindowMemory(k=2,return_messages=True,input_key="query")
memory_variable = memory.load_memory_variables({})
llm = ChatOpenAI(model="gpt-3.5-turbo")
chain = RunnablePassthrough.assign(history= RunnableLambda(memory.load_memory_variables) | itemgetter("history")) | prompt | llm |StrOutputParser()

while True:
    query = input("Human:")
    if query == "exit":
        exit(-1)
    chain_input = {"query": query, "history": []}
    response = chain.stream(chain_input)
    print("AI:", flush=True, end="")
    ai_output=""
    for chunk in response:
        print(chunk, flush=True, end="")
        ai_output += chunk
    memory.save_context(chain_input, {"output": ai_output})
    print("")
    print("history:", memory.load_memory_variables({}))

